import os
import argparse
import subprocess
import tempfile
import numpy as np
from pydub import AudioSegment
import noisereduce as nr

def extract_audio(input_mp4, output_audio):
    """使用FFmpeg从MP4中提取音频为WAV格式"""
    cmd = [
        'ffmpeg',
        '-y',  # 覆盖输出文件
        '-i', input_mp4,
        '-vn',  # 忽略视频
        '-acodec', 'pcm_s16le',  # PCM 16位小端格式
        '-ar', '44100',  # 采样率
        '-ac', '1',  # 单声道（降噪效果更好）
        output_audio
    ]
    subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

def denoise_audio(input_audio, output_audio, strength=0.85):
    """
    对音频进行降噪处理
    :param strength: 降噪强度 (0.0-1.0)
    """
    # 加载音频
    audio = AudioSegment.from_file(input_audio)
    
    # 转换为numpy数组
    samples = np.array(audio.get_array_of_samples())
    
    # 应用降噪
    reduced_noise = nr.reduce_noise(
        y=samples.astype(np.float32),
        sr=audio.frame_rate,
        stationary=True,
        prop_decrease=strength,
        n_fft=1024,
        win_length=1024,
        hop_length=256
    )
    
    # 创建处理后的音频
    denoised_audio = AudioSegment(
        reduced_noise.astype(np.int16).tobytes(),
        frame_rate=audio.frame_rate,
        sample_width=audio.sample_width,
        channels=1
    )
    
    # 保存处理后的音频
    denoised_audio.export(output_audio, format="wav")

def replace_audio(input_mp4, input_audio, output_mp4, audio_bitrate='192k'):
    """
    将处理后的音频替换回MP4文件
    :param audio_bitrate: 输出音频比特率
    """
    cmd = [
        'ffmpeg',
        '-y',  # 覆盖输出文件
        '-i', input_mp4,  # 原始视频
        '-i', input_audio,  # 处理后的音频
        '-c:v', 'copy',  # 直接复制视频流
        '-c:a', 'aac',  # 重新编码音频为AAC
        '-b:a', audio_bitrate,  # 音频比特率
        '-map', '0:v:0',  # 使用第一个输入的视频流
        '-map', '1:a:0',  # 使用第二个输入的音频流
        '-shortest',
        output_mp4
    ]
    subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

def process_video(input_path, output_path, strength=0.85, bitrate='192k'):
    """处理MP4文件：提取音频→降噪→替换音频"""
    with tempfile.TemporaryDirectory() as tmp_dir:
        # 创建临时文件路径
        original_audio = os.path.join(tmp_dir, "original_audio.wav")
        denoised_audio = os.path.join(tmp_dir, "denoised_audio.wav")
        
        # 步骤1: 提取音频
        print(f"提取音频: {input_path}")
        extract_audio(input_path, original_audio)
        
        # 步骤2: 降噪处理
        print(f"降噪处理 (强度: {strength})")
        denoise_audio(original_audio, denoised_audio, strength)
        
        # 步骤3: 替换音频并输出MP4
        print(f"合成视频: {output_path}")
        replace_audio(input_path, denoised_audio, output_path, bitrate)
        
        print(f"处理完成! 已保存至: {output_path}")

def main():
    # 设置命令行参数
    parser = argparse.ArgumentParser(description='MP4音频降噪处理工具')
    parser.add_argument('-i', '--input', required=True, help='输入MP4文件路径')
    parser.add_argument('-o', '--output', help='输出MP4文件路径（可选）')
    parser.add_argument('-s', '--strength', type=float, default=0.85, 
                        help='降噪强度 (0.0-1.0, 默认0.85)')
    parser.add_argument('-b', '--bitrate', default='192k', 
                        help='输出音频比特率 (默认192k)')
    
    args = parser.parse_args()
    
    # 处理输出路径
    if not args.output:
        # 自动生成输出文件名
        base, ext = os.path.splitext(args.input)
        args.output = f"{base}_denoised{ext}"
    
    # 验证参数
    if not os.path.exists(args.input):
        print(f"错误: 输入文件不存在 - {args.input}")
        return
    
    if not 0.0 <= args.strength <= 1.0:
        print(f"错误: 降噪强度必须在0.0到1.0之间")
        return
    
    # 处理视频
    try:
        process_video(args.input, args.output, args.strength, args.bitrate)
    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        print("请确保已安装FFmpeg并添加到系统PATH")

if __name__ == "__main__":
    main()